package top.sqdpt.rag.handler;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.content.retriever.neo4j.Neo4jGraph;
import dev.langchain4j.rag.content.retriever.neo4j.Neo4jText2CypherRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.springframework.stereotype.Component;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

@Component
public class KnowledgeGraphRetrievalHandler {

    private final ContentRetriever vectorRetriever;
    private final Neo4jText2CypherRetriever graphRetriever;
    private final ChatLanguageModel chatModel;

    public KnowledgeGraphRetrievalHandler(EmbeddingModel embeddingModel,
                                          EmbeddingStore<TextSegment> embeddingStore,
                                          Neo4jGraph neo4jGraph,
                                          ChatLanguageModel chatModel) {
        this.chatModel = chatModel;
        this.vectorRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(5)
                .minScore(0.618)
                .build();

        graphRetriever = new Neo4jText2CypherRetriever(neo4jGraph,chatModel,null);
    }


    public List<Content> retrieveHybrid(Query query) {
        // 1. 向量检索
        List<Content> vectorResults = vectorRetriever.retrieve(query);

        // 2. 知识图谱检索
        List<Content> graphResults = graphRetriever.retrieve(query);

        // 3. 合并结果
        List<Content> combinedResults = combineResults(vectorResults, graphResults);

        // 4. 重排序
        return rerankResults(query.text(), combinedResults);
    }

    private List<Content> combineResults(List<Content> vectorResults, List<Content> graphResults) {
        // 简单合并策略，可根据需要优化
        return Stream.concat(vectorResults.stream(), graphResults.stream())
                .distinct()
                .toList();
    }


    private List<Content> rerankResults(String query, List<Content> contents) {
        if (contents.size() <= 1) {
            return contents;
        }

        // 使用Cohere进行重排序
        PromptTemplate template = PromptTemplate.from("""
            请根据以下问题和文档的相关性进行排序:
            问题: {{query}}
            
            文档列表:
            {{#each contents}}
            - {{this.textSegment.text}}
            {{/each}}
            
            返回排序后的索引列表，例如 [2,0,1]
            """);

        Prompt prompt = template.apply(Map.of(
                "query", query,
                "contents", contents
        ));

        ChatResponse response = chatModel.chat(prompt.toUserMessage());
        String orderedIndices = response.aiMessage().text();

        // 解析排序结果并应用
        return parseAndApplyOrder(orderedIndices, contents);
    }

    private List<Content> parseAndApplyOrder(String orderedIndices, List<Content> contents) {
        // 实现解析逻辑
        // 简化为示例，实际应处理各种格式
        try {
            List<Integer> indices = Arrays.stream(orderedIndices
                            .replaceAll("[\\[\\]]", "")
                            .split(","))
                    .map(String::trim)
                    .map(Integer::parseInt)
                    .toList();

            return indices.stream()
                    .map(contents::get)
                    .toList();
        } catch (Exception e) {
            return contents; // 解析失败返回原始顺序
        }
    }




}
